import base64
import os
import json
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding as asym_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import requests

def load_server_public_key(public_key_path):
    """加载服务端RSA公钥"""
    with open(public_key_path, "rb") as key_file:
        public_key = serialization.load_pem_public_key(
            key_file.read(),
            backend=default_backend()
        )
    return public_key

def encrypt_data(server_public_key):
    """加密流程：RSA加密AES密钥 + AES-GCM加密数据（IV和Tag内嵌）"""
    # 1. 生成随机AES密钥（32字节，AES-256）
    aes_key = os.urandom(32)

    # 2. 生成随机IV（12字节，GCM标准长度）
    iv = os.urandom(12)

    # 3. 使用AES-GCM加密数据
    plaintext = b'{"title": "Python"}'
    encryptor = Cipher(
        algorithms.AES(aes_key),
        modes.GCM(iv),
        backend=default_backend()
    ).encryptor()
    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
    tag = encryptor.tag  # 获取GCM认证标签（16字节）
    # 4. 合并 IV + 密文 + Tag 为一个字节流
    encrypted_data = ciphertext + tag
    # 5. 使用RSA-OAEP加密AES密钥
    encrypted_aes_key = server_public_key.encrypt(
        aes_key,
        asym_padding.OAEP(
            mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

    return {
        "encrypted_aes_key": base64.b64encode(encrypted_aes_key).decode("utf-8"),
        "encrypted_data": base64.b64encode(encrypted_data).decode("utf-8"),
        "encrypted_aes_iv":base64.b64encode(iv).decode("utf-8"),
        "aes_key":base64.b64encode(aes_key).decode("utf-8"),
    }

def decryption(aes_key_bytes, encrypted_content, iv_base64):
    """
    使用AES-GCM解密数据。

    :param aes_key_bytes: AES密钥的字节串（16, 24 或 32 字节）
    :param encrypted_content: 要解密的密文字节串
    :param iv_base64: 初始化向量（IV）的Base64编码字符串
    :return: 解密后的明文字节串
    :raises Exception: 如果解密过程中出现错误
    """
    # 解码Base64 IV
    aes_key = base64.b64decode(aes_key_bytes)
    content = base64.b64decode(encrypted_content)
    decryContent = content[:-16]
    tag = content[-16:]
    try:
        iv = base64.b64decode(iv_base64)
    except Exception as e:
        raise ValueError("无效的IV Base64编码") from e

    # 验证IV长度（GCM推荐使用12字节的IV）
    if len(iv) != 12:
        raise ValueError("IV长度应为12字节")

    # 验证密钥长度（16, 24 或 32 字节对应AES-128, AES-192, AES-256）
    if len(aes_key) not in (16, 24, 32):
        raise ValueError("AES密钥长度必须为16, 24 或 32 字节")

    # 创建AES-GCM解密器
    try:
        decryptor = Cipher(
            algorithms.AES(aes_key),
            modes.GCM(iv, tag),  # java隐式tag 会将密文和tag放在一起返回，解密的时候需要分开
            backend=default_backend()
        ).decryptor()
    except Exception as e:
        raise ValueError("无效的AES密钥或IV") from e

    # 解密数据
    try:
        decrypted_padded = decryptor.update(decryContent) #+ decryptor.finalize()
    except Exception as e:
        raise ValueError("解密失败，可能是由于数据损坏或密钥错误") from e
    return decrypted_padded

def send_encrypted_request(encrypted_data):
    """发送加密请求到服务端"""
    headers = {
        "X-Encrypted-AES-Key": encrypted_data["encrypted_aes_key"],
        "X-Encrypted-AES-IV": encrypted_data["encrypted_aes_iv"],
        "Content-Type": "application/json"
    }
    response = requests.post(
        "http://localhost:8080/dataShare/testEncryption",
        headers=headers,
        data=encrypted_data["encrypted_data"]
    )
    return response.text,response.headers['X-Encrypted-AES-IV']

if __name__ == "__main__":
    # 加载服务端公钥（替换为你的公钥路径）
    server_public_key = load_server_public_key("../key/public_key.pem")

    # 加密数据
    encrypted_data = encrypt_data(server_public_key)

    # 发送请求
    response,iv = send_encrypted_request(encrypted_data)

    # 解密返回的内容
    body = decryption(encrypted_data['aes_key'], response, iv)
    print("body:" , body.decode("utf-8"))